package com.itheima.ai.tools;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.List;

@Component
@RequiredArgsConstructor
public class VectorTools {


    private final VectorStore vectorStore;
    public void WritePDFtoVectorDatabase(Resource resource){
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource,
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder().build())
                        .withPagesPerDocument(1)
                        .build()
        );
        List<Document> read = reader.read();
        try {
            vectorStore.add(read);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public List<Document> SimilaritySearch(String context){
        List<Document> documents = vectorStore.similaritySearch(context);
        if (documents == null || documents.isEmpty()) {
            return List.of();
        }
        return documents;
    }

    public List<Document> SimilaritySearch(SearchRequest request){
        List<Document> documents = vectorStore.similaritySearch(request);
        if (documents == null || documents.isEmpty()) {
            return List.of();
        }
        return documents;
    }
}
